import re
import asyncio
import websockets


callbacks = []


# 解析请求参数: {"type": "request", "data": "hello,world!"}
async def parse(message):
    ret = {}
    message = message[1:-1]
    for item in [message[:message.index(",")], message[message.index(",") + 1:]]:
        key = item[:item.index(":")]
        value = item[item.index(":") + 1:]
        ret[key[1:-1]] = value[1:-1]

    return ret


async def auto_login(message, websocket):
    message = await parse(message)
    if message.get("type") == "response" and \
            re.search(r"^\\r\\u001b\[K\[xxxx@[A-Z,a-z,0-9,-]+\s~\]\$ $", message.get("data")):
        await websocket.send(str({"type": "request", "data": "su\r"}))
    elif message.get("type") == "response" and message.get("data") == "Password: ":
        await websocket.send(str({"type": "request", "data": "xxxx\r"}))
    else:
        return False

    return True


async def echo(websocket, path):
    #fetch msg
    async for message in websocket:
        print(f"message{websocket.remote_address}: {message}")
        flag = True
        for callback in callbacks:
            if await callback(message, websocket):
                flag = False

        if flag:
            await websocket.send(str({"data": "ok"}))


async def main():
    # start a websocket server
    async with websockets.serve(echo, "0.0.0.0", 61819):
        await asyncio.Future()  # run forever


# 回调函数集
callbacks.append(auto_login)
asyncio.run(main())
